import asyncio
import json
import logging
import gzip
import base64

import tornado
import tornado.ioloop
import tornado.websocket
from bson.objectid import ObjectId

from api_base import *
from docker_client import *
from runner import my_run_on_executor


class WebSocketContainerFileHandler(tornado.websocket.WebSocketHandler):
    def check_origin(self, origin: str):
        return True

    def get_compression_options(self):
        return {
            'compression_level': 6,
            'mem_level': 6
        }

    async def on_message(self, message):
        try:
            js = json.loads(message)
            op = js['op']
            filename = js['filename']
            if op == 'put':
                payload = js['payload']
        except:
            self.write_message({'info': 'json_error'})
            return
        containerid = get_id(self.request.uri)
        if op == 'get':
            stream = await self.get_file_from_docker(containerid, filename)
            try:
                for chunk in stream:
                    await self.write_message(chunk, binary=True)
                self.close(reason='EOF')
            except:
                pass
        elif op == 'put':
            await self.put_file_to_docker(containerid, filename, payload)
        else:
            self.write_message({'info': 'unknown op'})
            return

    @my_run_on_executor
    def get_file_from_docker(self, containerid, filename):
        try:
            container = client.containers.get(containerid)
            stream, _ = container.get_archive(filename)
            return stream
        except:
            self.write_message({'info': 'fail'})
            return []

    @my_run_on_executor
    def put_file_to_docker(self, containerid, filepath, payload):
        payload = gzip.decompress(base64.decodebytes(bytes(payload)))
        container = client.containers.get(containerid)
        container.put_archive(filepath, payload)


async def recv_from_container(reader, ws):
    while True:
        data = await reader.read(1024)
        if not data:
            logging.info('container closed')
            break
        try:
            ws.write_message(data)
        except tornado.websocket.WebSocketClosedError:
            logging.info('ws closed')
            break


class WebSocketTerminalHandler(tornado.websocket.WebSocketHandler):
    sockets = {}
    sock_reader = {}
    sock_writer = {}

    async def open(self):
        logging.info('open: {}'.format(self.request.uri))
        containerid = self.request.uri.split('/')[-1]
        try:
            container = await get_container(containerid)
            container.start()
            self.sockets[containerid] = container.attach_socket(
                params={'stdin': 1, 'stdout': 1, 'stderr': 1, 'stream': 1})
            self.sock_reader[containerid], self.sock_writer[containerid] =\
                await asyncio.open_connection(sock=self.sockets[containerid]._sock)
            loop = tornado.ioloop.IOLoop.current()
            loop.add_callback(recv_from_container,
                              self.sock_reader[containerid], self)
            self.sock_writer[containerid].write('\n'.encode())
        except:
            logging.error('Can\'t attach socket to container '+containerid)
            self.close(
                code=404, reason='Can\'t attach socket to container '+containerid)
        self.set_nodelay(True)

    async def on_message(self, message):
        containerid = self.request.uri.split('/')[-1]
        self.sock_writer[containerid].write(message.encode())

    def check_origin(self, origin: str):
        return True

    def on_close(self):
        containerid = self.request.uri.split('/')[-1]
        try:
            self.sock_reader[containerid] = None
            self.sock_writer[containerid] = None
            self.sockets[containerid].close()
            self.sockets[containerid] = None
        except:
            logging.error('websocket on_close')
